import com.alibaba.fastjson.JSONObject;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import org.apache.calcite.adapter.jdbc.JdbcSchema;
import org.apache.calcite.config.CalciteConnectionConfigImpl;
import org.apache.calcite.interpreter.Bindables;
import org.apache.calcite.jdbc.CalciteConnection;
import org.apache.calcite.jdbc.CalciteSchema;
import org.apache.calcite.plan.RelOptTable;
import org.apache.calcite.plan.hep.HepPlanner;
import org.apache.calcite.plan.hep.HepProgram;
import org.apache.calcite.prepare.CalciteCatalogReader;
import org.apache.calcite.rel.RelHomogeneousShuttle;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.RelRoot;
import org.apache.calcite.rel.RelShuttle;
import org.apache.calcite.rel.core.TableScan;
import org.apache.calcite.rel.logical.LogicalTableScan;
import org.apache.calcite.rel.type.RelDataTypeSystem;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.sql.SqlNode;
import org.apache.calcite.sql.fun.SqlStdOperatorTable;
import org.apache.calcite.sql.parser.SqlParser;
import org.apache.calcite.sql.type.SqlTypeFactoryImpl;
import org.apache.calcite.sql.validate.SqlValidator;
import org.apache.calcite.sql.validate.SqlValidatorUtil;
import org.apache.calcite.tools.FrameworkConfig;
import org.apache.calcite.tools.Frameworks;
import org.apache.calcite.tools.RelBuilder;
import org.apache.calcite.tools.RelRunner;

import javax.sql.DataSource;
import java.sql.*;
import java.util.List;
import java.util.Map;
import java.util.Properties;

public class CalciteJdbcExample {

    private static final String POSTGRESQL_SCHEMA = "PUBLIC";
    private static final String MYSQL_SCHEMA = "mysql";

    public static void main(String[] args) throws Exception {

        // Build our connection
        Connection connection = DriverManager.getConnection("jdbc:calcite:");

        // Unwrap our connection using the CalciteConnection
        CalciteConnection calciteConnection = connection.unwrap(CalciteConnection.class);

        // Get a pointer to our root schema for our Calcite Connection
        SchemaPlus rootSchema = calciteConnection.getRootSchema();

        // Instantiate a data source, this can be autowired in using Spring as well
        DataSource postgresDataSource = JdbcSchema.dataSource(
                "jdbc:postgresql://localhost/db",
                "org.postgresql.Driver", // Change this if you want to use something like MySQL, Oracle, etc.
                "postgres", // username
                "example"   // password
        );

        // Instantiate a data source, this can be autowired in using Spring as well
        DataSource mysqlDataSource = JdbcSchema.dataSource(
                "jdbc:mysql://192.168.98.228:3306/skypay",
                "com.mysql.jdbc.Driver", // Change this if you want to use something like MySQL, Oracle, etc.
                "pay", // username
                "querypay!321"   // password
        );

        // Attach our Postgres Jdbc Datasource to our Root Schema
        //rootSchema.add(POSTGRESQL_SCHEMA, JdbcSchema.create(rootSchema, POSTGRESQL_SCHEMA, postgresDataSource, null, null));

        // Attach our MySQL Jdbc Datasource to our Root Schema
        rootSchema.add(MYSQL_SCHEMA, JdbcSchema.create(rootSchema, MYSQL_SCHEMA, mysqlDataSource, null, null));


        // Build a framework config to attach to our Calcite Planners and  Optimizers
        FrameworkConfig config = Frameworks.newConfigBuilder()
                .defaultSchema(rootSchema)
                .build();

        RelBuilder rb = RelBuilder.create(config);

        RelNode node = rb
                // First parameter is the Schema, the second is the table name
                .scan("PUBLIC", "TABLE_NAME_IN_POSTGRES")
                .scan("mysql", "TABLE_NAME_IN_MYSQL")
                // If you want to select from more than one table, you can do so by adding a second scan parameter
                .filter(
                        rb.equals(rb.field("fieldname"), rb.literal("literal"))
                )
                // These are the fields you want to return from your query
                .project(
                        rb.field("id"),
                        rb.field("col1"),
                        rb.field("colb")
                )
                .build();

        String sql = "select u.id as user_id, u.name as user_name, j.company as user_company, u.age as user_age " +
                " from users u join jobs j on u.name=j.name " +
                " where u.age > 30 and j.id>10 " +
                " order by user_id";

        SqlParser.Config sqlConfig = SqlParser.configBuilder().build();
        SqlParser parser = SqlParser.create(sql, sqlConfig);
        SqlNode sqlNode = parser.parseQuery();

        //note: 二、sql validate（会先通过Catalog读取获取相应的metadata和namespace）
//note: get metadata and namespace
        SqlTypeFactoryImpl factory = new SqlTypeFactoryImpl(RelDataTypeSystem.DEFAULT);
        CalciteCatalogReader calciteCatalogReader = new CalciteCatalogReader(
                CalciteSchema.from(rootSchema),
                CalciteSchema.from(rootSchema).path(null),
                factory,
                new CalciteConnectionConfigImpl(new Properties()));

//note: 校验（包括对表名，字段名，函数名，字段类型的校验。）
      /*  SqlValidator validator = SqlValidatorUtil.newValidator(SqlStdOperatorTable.instance(), calciteCatalogReader, factory,
                conformance(frameworkConfig));
        SqlNode validateSqlNode = validator.validate(sqlNode);*/

        //step1、SQL 解析阶段（SQL–>SqlNode）
        //SqlParser parser = SqlParser.create(sql, SqlParser.Config.DEFAULT);



        HepProgram program = HepProgram.builder().build();
        HepPlanner planner = new HepPlanner(program);
        planner.setRoot(node);

        RelNode optimizedNode = planner.findBestExp();

        final RelShuttle shuttle = new RelHomogeneousShuttle() {
            @Override public RelNode visit(TableScan scan) {
                final RelOptTable table = scan.getTable();
                if (scan instanceof LogicalTableScan && Bindables.BindableTableScan.canHandle(table)) {
                    return Bindables.BindableTableScan.create(scan.getCluster(), table);
                }
                return super.visit(scan);
            }
        };

        optimizedNode = optimizedNode.accept(shuttle);

        final RelRunner runner = connection.unwrap(RelRunner.class);
        PreparedStatement ps = runner.prepare(optimizedNode);

        ps.execute();

        ResultSet resultSet = ps.getResultSet();
        //DBTablePrinter.printResultSet(resultSet);
        System.out.println(JSONObject.toJSONString(getData(resultSet)));
    }

    public static List<Map<String,Object>> getData(ResultSet resultSet)throws Exception{
        List<Map<String,Object>> list = Lists.newArrayList();
        ResultSetMetaData metaData = resultSet.getMetaData();
        int columnSize = metaData.getColumnCount();

        while (resultSet.next()) {
            Map<String, Object> map = Maps.newLinkedHashMap();
            for (int i = 1; i < columnSize + 1; i++) {

                map.put(metaData.getColumnLabel(i).toLowerCase(), resultSet.getObject(i));
            }
            list.add(map);
        }
        return list;
    }
}